# server.py
import socket
import os
from cryptography.fernet import Fernet

KEY = b'2DArf8mfLdgsDgwdnBQ5Uzr9zIRSEVerCnVk8Mn9xTQ='
BUFFER_SIZE = 4096
SEPARATOR = "<SEPARATOR>"
SAVE_PATH = "received_files"

# 创建保存目录
os.makedirs(SAVE_PATH, exist_ok=True)

server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
server_socket.bind(("127.0.0.1", 8080))
server_socket.listen(5)
cipher = Fernet(KEY)

print("等待连接...")
client_socket, address = server_socket.accept()
print(f"收到来自 {address[0]}:{address[1]} 的连接")


def save_file(file_name, file_size):
    save_path = os.path.join(SAVE_PATH, file_name)
    bytes_received = 0

    with open(save_path, "wb") as f:
        while bytes_received < file_size:
            encrypted_data = client_socket.recv(BUFFER_SIZE)
            decrypted_data = cipher.decrypt(encrypted_data)
            f.write(decrypted_data)
            bytes_received += len(decrypted_data)
    print(f"文件 {file_name} 保存成功 ({bytes_received} bytes)")


while True:
    encrypted_data = client_socket.recv(BUFFER_SIZE)
    if not encrypted_data:
        break

    try:
        decrypted_data = cipher.decrypt(encrypted_data).decode("utf-8")
    except:
        print("解密失败！连接已终止")
        break

    # 文件传输处理
    if decrypted_data.startswith("FILE"):
        _, file_name, file_size = decrypted_data.split(SEPARATOR)
        file_size = int(file_size)
        print(f"正在接收文件: {file_name} ({file_size} bytes)")
        save_file(file_name, file_size)
        continue

    # 普通消息处理
    print(f"收到消息：{decrypted_data}")

    if decrypted_data == "exit":
        break

    reply = input("输入回复：")
    encrypted_reply = cipher.encrypt(reply.encode("utf-8"))
    client_socket.send(encrypted_reply)

client_socket.close()